from paddle import fluid
import numpy as np
from collections import OrderedDict
import pickle
import os
import shutil

class Test:
    def test(self):
        pass


class LeakyRelu(fluid.dygraph.Layer):
    def __init__(self, alpha=0.02):
        super(LeakyRelu, self).__init__()
        self.alpha = alpha

    def forward(self, x):
        return fluid.layers.leaky_relu(x, alpha=self.alpha)


class Model(fluid.dygraph.Layer):
    def __init__(self):
        super(Model, self).__init__()
        self.conv = fluid.dygraph.Conv2D(3, 4, 3, 1, 1)
        self.leaky_relu = LeakyRelu(alpha=0.2)

    def forward(self, x):
        x = self.conv(x)
        x = self.leaky_relu(x)
        return x


class LossAlignTest(Test):
    def __init__(self):
        pass

    def dump_params(self, model):
        for name, param in model.state_dict().items():
            print(f'{name}: {param.numpy()}')

    def save_params(self, model):
        params = OrderedDict()
        for name, param in model.state_dict().items():
            params[name] = param.numpy()
        pickle.dump(params, open('test_params.pkl', 'wb'))

    def load_params(self, model: Model):
        params = pickle.load(open('test_params.pkl', 'rb'))
        model.load_dict(params)

    def test(self):
        with fluid.dygraph.guard():
            self.model = Model()
            self.load_params(self.model)
            self.dump_params(self.model)
            self.optmizes = fluid.optimizer.AdamOptimizer(learning_rate=0.01,
                                                          parameter_list=self.model.parameters(),
                                                          beta1=0.0,
                                                          beta2=0.99,
                                                          regularization=fluid.regularizer.L2Decay(1e-4)
                                                          )
            for step in range(100):
                self.inputs = np.random.randn(1, 3, 3, 3).astype(np.float32)
                self.inputs.fill(0.5)
                self.inputs = fluid.dygraph.to_variable(self.inputs)
                self.y = fluid.dygraph.to_variable(np.ones((1, 4, 3, 3)).astype(np.float32))
                outputs = self.model(self.inputs)
                # print('\nstep:', step)
                # print('outputs:', outputs.numpy())
                loss = fluid.layers.reduce_mean(fluid.layers.abs(self.y - outputs))
                print('loss:', loss.numpy())
                loss.backward()
                self.optmizes.minimize(loss)
                self.model.clear_gradients()
            # self.save_params(self.model)


if __name__ == '__main__':
    test = LossAlignTest()

    test.test()
